use crate::{Allocator, Error, PoolAlloc, prelude::*};
use core::alloc::Layout;
use core::cell::Cell;
use core::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
use core::fmt::{self, Debug, Display, Formatter};
use core::hash::{Hash, Hasher};
use core::marker::PhantomData;
use core::mem::{ManuallyDrop, MaybeUninit};
use core::ops::Deref;
use core::ptr::{self, NonNull};

#[repr(C)]
struct RcInner<T, A: Allocator> {
    cnt: Cell<isize>,
    val: T,
    layout: Layout,
    alloc: ManuallyDrop<A>,
}

#[repr(C)]
pub struct Rc<'a, T, A: Allocator = PoolAlloc> {
    inner: NonNull<RcInner<T, A>>,
    mark1: PhantomData<RcInner<T, A>>,
    mark2: PhantomData<&'a A>,
}

impl Rc<'_, (), PoolAlloc> {
    pub fn init_then<T, F>(f: F) -> Result<Rc<'static, T, PoolAlloc>, Error> 
    where
        F: FnOnce(&mut MaybeUninit<T>) -> Result<(), Error>,
    {
        Rc::new_then(f)
    }
    pub fn init_with_buf<T, F>(layout: Layout, f: F) -> Result<Rc<'static, T, PoolAlloc>, Error> 
    where
        F: FnOnce(&mut MaybeUninit<T>, NonNull<[u8]>) -> Result<(), Error>,
    {
        Rc::new_with_buf(layout, f)
    }
}

impl<A: Allocator + Clone> Rc<'_, (), A> {
    pub fn init_then_in<T, F>(alloc: A, f: F) -> Result<Rc<'static, T, A>, Error> 
    where
        F: FnOnce(&mut MaybeUninit<T>) -> Result<(), Error>,
    {
        Rc::new_then_in(alloc, f)
    }
    pub fn init_with_buf_in<T, F>(alloc: A, layout: Layout, f: F) -> Result<Rc<'static, T, A>, Error> 
    where
        F: FnOnce(&mut MaybeUninit<T>, NonNull<[u8]>) -> Result<(), Error>,
    {
        Rc::new_with_buf_in(alloc, layout, f)
    }
}

impl<T> Rc<'static, T, PoolAlloc> {
    pub fn new(val: T) -> Result<Self, Error> {
        Self::new_in(PoolAlloc, val)
    }

    pub fn new_then<F>(f: F) -> Result<Self, Error>
    where
        F: FnOnce(&mut MaybeUninit<T>) -> Result<(), Error>,
    {
        Self::new_then_in(PoolAlloc, f)
    }

    pub fn new_with_buf<F>(layout: Layout, f: F) -> Result<Self, Error>
    where
        F: FnOnce(&mut MaybeUninit<T>, NonNull<[u8]>) -> Result<(), Error>,
    {
        Self::new_with_buf_in(PoolAlloc, layout, f)
    }
}

impl<T, A: Allocator + Clone> Rc<'_, T, A> {
    pub fn new_in(alloc: A, val: T) -> Result<Self, Error> {
        Self::new_then_in(alloc, |uninit| {
            uninit.write(val);
            Ok(())
        })
    }

    pub fn new_then_in<F>(alloc: A, f: F) -> Result<Self, Error>
    where
        F: FnOnce(&mut MaybeUninit<T>) -> Result<(), Error>,
    {
        unsafe {
            alloc
                .alloc_then(|inner: &mut MaybeUninit<RcInner<T, A>>| {
                    let this = inner.as_mut_ptr().cast::<RcInner<T, A>>();
                    ptr::addr_of_mut!((*this).cnt).write(Cell::new(1));
                    ptr::addr_of_mut!((*this).alloc).write(ManuallyDrop::new(alloc.clone()));
                    ptr::addr_of_mut!((*this).layout).write(Layout::new::<RcInner<T, A>>());
                    f(&mut *ptr::addr_of_mut!((*this).val).cast::<MaybeUninit<T>>())
                })
                .map(|inner| Self {
                    inner,
                    mark1: PhantomData,
                    mark2: PhantomData,
                })
        }
    }
    pub fn new_with_buf_in<F>(alloc: A, layout: Layout, f: F) -> Result<Self, Error>
    where
        F: FnOnce(&mut MaybeUninit<T>, NonNull<[u8]>) -> Result<(), Error>,
    {
        let (all_layout, offset) = Layout::new::<RcInner<T, A>>()
            .extend(layout)
            .map_err(|_| Error::new(EINVAL))?;
        unsafe {
            alloc
                .alloc_buf(all_layout, |ptr| {
                    let this = ptr.cast::<RcInner<T, A>>().as_ptr();
                    let buf = NonNull::from(&ptr.as_ref()[offset..]);
                    ptr::addr_of_mut!((*this).cnt).write(Cell::new(1));
                    ptr::addr_of_mut!((*this).alloc).write(ManuallyDrop::new(alloc.clone()));
                    ptr::addr_of_mut!((*this).layout).write(all_layout);
                    f(
                        &mut *ptr::addr_of_mut!((*this).val).cast::<MaybeUninit<T>>(),
                        buf,
                    )
                })
                .map(|ptr| Self {
                    inner: ptr.cast::<RcInner<T, A>>(),
                    mark1: PhantomData,
                    mark2: PhantomData,
                })
        }
    }
}

impl<T, A> Drop for Rc<'_, T, A>
where
    A: Allocator,
{
    fn drop(&mut self) {
        let cnt = self.get_inner().cnt.get();
        if cnt == 1 {
            let inner = unsafe { self.inner.as_mut() };
            let alloc = unsafe { ManuallyDrop::take(&mut inner.alloc) };
            let layout = inner.layout;
            unsafe { ptr::drop_in_place(inner) };
            unsafe { alloc.release_with(self.inner, layout) };
        } else {
            self.get_inner().cnt.set(cnt - 1);
        }
    }
}

impl<T, A> Rc<'_, T, A>
where
    A: Allocator,
{
    pub fn get_mut(&mut self) -> Option<&mut T> {
        if self.get_inner().cnt.get() == 1 {
            Some(unsafe { self.get_mut_unchecked() })
        } else {
            None
        }
    }

    /// # Safety
    /// Rc使用者保证可写借用和只读借用并存，或者多个可写借用是安全的
    /// Rc不支持Send/Sync，因此无需考虑并发冲突的风险
    pub unsafe fn get_mut_unchecked(&mut self) -> &mut T {
        unsafe { &mut (*self.inner.as_ptr()).val }
    }

    fn get_inner(&self) -> &RcInner<T, A> {
        unsafe { self.inner.as_ref() }
    }
}

impl<T, A: Allocator> AsRef<T> for Rc<'_, T, A> {
    #[inline(always)]
    fn as_ref(&self) -> &T {
        &self.get_inner().val
    }
}

impl<T, A: Allocator> Clone for Rc<'_, T, A> {
    fn clone(&self) -> Self {
        let inner = self.get_inner();
        inner.cnt.set(inner.cnt.get() + 1);
        Self {
            inner: self.inner,
            mark1: PhantomData,
            mark2: PhantomData,
        }
    }
}

impl<T, A: Allocator> Unpin for Rc<'_, T, A> {}

impl<T, A: Allocator> Deref for Rc<'_, T, A> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<T: Display, A: Allocator> Display for Rc<'_, T, A> {
    #[inline]
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        Display::fmt(&**self, f)
    }
}

impl<T: Debug, A: Allocator> Debug for Rc<'_, T, A> {
    #[inline]
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        Debug::fmt(&**self, f)
    }
}

impl<T, A: Allocator> fmt::Pointer for Rc<'_, T, A> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ptr = self.as_ref() as *const _;
        fmt::Pointer::fmt(&ptr, f)
    }
}

impl<T: PartialEq, A: Allocator> PartialEq for Rc<'_, T, A> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(&**self, &**other)
    }
}

impl<T: PartialOrd, A: Allocator> PartialOrd for Rc<'_, T, A> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
    #[inline]
    fn lt(&self, other: &Self) -> bool {
        PartialOrd::lt(&**self, &**other)
    }
    #[inline]
    fn le(&self, other: &Self) -> bool {
        PartialOrd::le(&**self, &**other)
    }
    #[inline]
    fn gt(&self, other: &Self) -> bool {
        PartialOrd::gt(&**self, &**other)
    }
    #[inline]
    fn ge(&self, other: &Self) -> bool {
        PartialOrd::ge(&**self, &**other)
    }
}

impl<T: Ord, A: Allocator> Ord for Rc<'_, T, A> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: Eq, A: Allocator> Eq for Rc<'_, T, A> {}

impl<T: Hash, A: Allocator> Hash for Rc<'_, T, A> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        Hash::hash(&**self, state)
    }
}

#[cfg(test)]
mod test {
    use crate::{MemPool, Rc};
    #[test]
    fn test() {
        let rc = Rc::new(100).unwrap();

        let cnt = *rc + 100;
        assert_eq!(200, cnt);

        let mut rc2 = rc.clone();
        let ret = rc2.get_mut();
        assert_eq!(ret, None);
    }

    #[test]
    fn test_pool() {
        static mut DROP: usize = 0;
        struct Foo;
        impl Drop for Foo {
            fn drop(&mut self) {
                unsafe { DROP += 1 };
            }
        }
        let pool = MemPool::new_boxed(0).unwrap().leak();
        unsafe { DROP = 0 };
        {
            let rc = Rc::new_in(&*pool, Foo);
            {
                let _rc1 = rc.clone();
            }
            unsafe { assert_eq!(DROP, 0) };
        }
        unsafe { assert_eq!(DROP, 1) };
    }
}
